import gradio as gr
import numpy as np
from PIL import Image
import sys
import subprocess
from pathlib import Path
import pinecone
from pinecone import Pinecone, ServerlessSpec
from collections import Counter

PINECONE_API_KEY = "pcsk_4tZGoM_HLqw5naiHZaXJWLh4WY3W3MEEcozbMT5Ci3cUkYpTGSJXNdLZZWzUQyZcyp9MVp"
INDEX_NAME = "mnist-index"
K = 11

BASE_DIR = Path(__file__).parent
TRAINING_SCRIPT = BASE_DIR / "pinecone_train.py"

def init_pinecone():
    pc = Pinecone(api_key=PINECONE_API_KEY)
    
    if INDEX_NAME not in pc.list_indexes().names():
        print(f"Pinecone index '{INDEX_NAME}' does not exist, running training script to create...")
        if not TRAINING_SCRIPT.exists():
            raise FileNotFoundError(f"Training script not found: {TRAINING_SCRIPT}, please create the index first")
        
        completed = subprocess.run([sys.executable, str(TRAINING_SCRIPT)], cwd=str(BASE_DIR))
        if completed.returncode != 0:
            raise RuntimeError(f"Training script execution failed, return code: {completed.returncode}")
    
    return pc.Index(INDEX_NAME)

index = init_pinecone()

def predict(my_dict):
    a = my_dict['composite']
    a = np.array(a)
    a = a[:, :, 3]
    a = 255 - a
    b = Image.fromarray(a, 'L')
    c = b.resize((8, 8))
    arr = np.array(c) / 255.0 * 16
    arr = (arr > 8) * 16
    Image.fromarray(arr.astype(np.uint8), 'L').save('input_debug.png')
    query_vector = arr.ravel().tolist()
    
    results = index.query(
        vector=query_vector,
        top_k=K,
        include_metadata=True
    )
    
    if not results['matches']:
        raise ValueError("Pinecone returned no search results, please check index data")
    
    labels = [match['metadata']['label'] for match in results['matches']]
    predicted_label = Counter(labels).most_common(1)[0][0]
    
    return int(predicted_label)

iface = gr.Interface(
    fn=predict,
    inputs=gr.Sketchpad(),
    outputs=gr.Label(num_top_classes=1),
    title="KNN Handwritten Digit Recognition (Pinecone Cloud Version)",
    description="KNN handwritten digit recognition based on Pinecone cloud vector database (each inference calls the cloud index)"
)

print("Launching Gradio in local mode...")
iface.launch(share=False) 